import sys
import json
from llama_index.llms.ollama import Ollama
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.embeddings.fastembed import FastEmbedEmbedding
from llama_index.core import Settings
from llama_index.core.query_engine import NLSQLTableQueryEngine
import openai
import logging
from llama_index.llms.openai import OpenAI
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
import uuid
from bs4 import BeautifulSoup
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.callbacks import CallbackManager
from llama_index.core.agent import AgentRunner
from llama_index.agent.openai import OpenAIAgent
from llama_index.agent.openai import OpenAIAgentWorker
from llama_index.embeddings.ollama import OllamaEmbedding
'''
pip install beautifulsoup4
pip install chromadb
pip install llama-index
pip install llama-index-llms-ollama
pip install llama-index-embeddings-fastembed
pip install llama_index.vector_stores.chroma
pip install llama-index-llms-openai
'''


# from actions.api.mysql_tab import MysqlTab
logger = logging.getLogger(__name__)

class RagUI:
    """ RagUI """

    def __init__(self ):
        logger.info(' RagUI -init')
        # openai.api_key = ''
        # openai.base_url = ''

        self.llm = OpenAI(temperature=0.19, model="gpt-4o")
        #self.llm = Ollama(model="llama3:8b-instruct-q2_K", base_url="http://localhost:9090",request_timeout=60000.0)
        # self.embed_model = FastEmbedEmbedding(model_name="BAAI/bge-base-en-v1.5")   
        ''' self.embed_model = OllamaEmbedding(
            model_name="llama3:8b-instruct-q2_K",
            base_url="http://localhost:9090",
            ollama_additional_kwargs={"mirostat": 0},
        )'''
        Settings.llm = self.llm
        # Settings.embed_model = self.embed_model
        self.callback_manager = CallbackManager([])
    
    def agent_ui(self,doc_type,query_str):
        # doc_type = doc-ui
        db = chromadb.PersistentClient(path=r"E:\100rag\chroma")
        chroma_collection = db.get_or_create_collection(doc_type)
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        index = VectorStoreIndex.from_vector_store(
            vector_store,
            embed_model=self.embed_model,
        )
         # query_engine = index.as_query_engine(similarity_top_k=2, llm=gpt35_llm)
        doc_engine = index.as_query_engine(similarity_top_k=6)
        
        # query
        prompt=f'''实现功能如下:
1:把xml转成json字符串.
2:请只回复答案，不要写解释.
请替换下面字符串:{query_str}
        '''
        
        # define tools
        query_engine_tools = []
        query_engine_tools.append(
            QueryEngineTool(
                query_engine=doc_engine,
                metadata=ToolMetadata(
                    name=f"vector_tool_{doc_type}",
                    description=(
                        "Useful for questions related to specific aspects of: 替换字符串"
                    ),
                ),
            )
        )
        # 创建agent
        openai_step_engine = OpenAIAgentWorker.from_tools(
            query_engine_tools, llm=self.llm, verbose=False
        )
        # 运行agent
        agent = AgentRunner(openai_step_engine)
        agt_res = agent.chat(prompt)
        if hasattr(agt_res, 'response'):
            return agt_res.response
        return '无法识别'
    
    def text2html(self,query_str: str):
        prompt=f'''
        按要求生成html5页面,语言是中文。        
        要求是:{query_str}
        请只回复答案，不要写解释。
        '''
        html_str = self.llm.complete(prompt)
        print(html_str.text)
        html = html_str.text
        # 假设这是你的HTML字符串
        return html
    
    def rag(self,query_str: str):    
        div = ""
        # 使用BeautifulSoup加载HTML字符串
        soup = BeautifulSoup(query_str, 'html.parser')
        # 找到<body>标签
        body = soup.find('body')
        # 遍历<body>标签的所有子对象
        for child in body.children:
            # 检查child是否是一个Tag对象
            if child.name is  None:
                continue
            # 打印子对象的完整内部HTML
            print(f"Tag: {child.name}, InnerHTML: {str(child)}")
            div += self.readDiv(child)
        new_uuid = str(uuid.uuid4())
        page ='{"id":"{'+new_uuid+'}","style":{"padding-left":"20px","padding-right":"20px","padding-bottom":"20px"},"FolderId":null,"ParentPageId":null,"children":['+div+'],"containerLeavel":0,"name":"Page","label":"page","type":"page","code":"apj_page","moduleId":"811C0F4CEA8E"}' # type: ignore
        return page
    
    def readDiv(self,divObj):
        if divObj is None:
            return ""
        div = ""
        for child in divObj:
            # 检查child是否是一个Tag对象
            if child.name is  None:
                continue
            # 打印子对象的完整内部HTML
            print(f"Tag: {child.name}, InnerHTML: {str(child)}")
            lw_name = child.name.lower()
            if lw_name == 'div':
                if '</div>' in str(child):
                    div += self.readDiv(child)
                else:
                    div += self.readUIDoc(str(child))
            if lw_name == 'input':
                div += self.readUIDoc(str(child))
            if lw_name == 'button':
                div += self.readUIDoc(str(child))
            if lw_name == 'form':
                div += self.readDiv(child)
        return div
        
    def readUIDoc(self,query_str): 
        db = chromadb.PersistentClient(path=r"E:\100rag\chroma")
        chroma_collection = db.get_or_create_collection("doc-ui")
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        index = VectorStoreIndex.from_vector_store(
            vector_store,
            embed_model=self.embed_model,
        )
         # query_engine = index.as_query_engine(similarity_top_k=2, llm=gpt35_llm)
        doc_engine = index.as_query_engine(similarity_top_k=6)
        
        # query
        prompt=f'''
Replace patterns based on contextual examples:
1:Replace XML with JSON.
2:Verify the JSON syntax. If the grammar is incorrect, Please replace the XML with JSON again.
3:Attributes such as "code", "codeType", "type", do not replace, keep unchanged.
4:Attributes such as "EN_US","ZH_CN" , "ZH_HK" should be replaced correctly,Content cannot be empty.
5:The validation attribute "ZH-CN" cannot be empty. If it is empty, please regenerate it.
6:Translate the content of the "EN_US" attribute into English.
7:don't write explanation.

XML={query_str}
JSON=?
'''
        rag_doc = doc_engine.query(prompt)
        print('#########################')
        # print(rag_doc)
        if hasattr(rag_doc, 'response'):
            return rag_doc.response
        return '无法识别'
        

ragUI = RagUI()
# ret = ragUI.agent_ui('doc-ui','<button class="px-4 py-2 bg-blue-500 text-white flex items-center"><i class="fas fa-camera mr-2"></i>提交</button>')
# print(ret)

# ret = ragUI.readUIDoc('<label class="block text-gray-700">你好</label>')
# print(ret)
#ret = ragUI.readUIDoc('<select class="mt-1 block w-full border border-gray-300 rounded-md p-2"><option>请选择客户来源</option></select>')
#print(ret)
#ret = ragUI.readUIDoc('<input type="text" id="cargo" name="cargo" placeholder="请输入商品">')
#print(ret)
#ret = ragUI.readUIDoc('<div class="flex flex-col items-center"></div>')
#print(ret)
html='''
<html>
<head>
    <script src="https://cdn.tailwindcss.com"></script>
</head>
<body class="bg-gray-100">
    <div class="max-w-4xl mx-auto mt-10 bg-white p-8 rounded-lg shadow-lg">
        <div class="flex justify-between items-center mb-6">
            <h2 class="text-xl font-semibold">新增</h2>
            <button class="text-gray-500">
                <i class="fas fa-expand"></i>
            </button>
        </div>
        <form>
            <div class="grid grid-cols-4 gap-6">
                <div class="col-span-1">
                    <label class="block text-gray-700">* 线索名称</label>
                    <input type="text" placeholder="请输入线索名称" class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">客户来源</label>
                    <select class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                        <option>请选择客户来源</option>
                    </select>
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">手机</label>
                    <input type="text" placeholder="请输入手机" class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">* 负责人</label>
                    <select class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                        <option>微脑管理员</option>
                    </select>
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">电话</label>
                    <input type="text" placeholder="请输入电话" class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">邮箱</label>
                    <input type="text" placeholder="请输入邮箱" class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">微信</label>
                    <input type="text" placeholder="请输入微信" class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">QQ</label>
                    <input type="text" placeholder="请输入 QQ" class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">客户行业</label>
                    <select class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                        <option>请选择客户行业</option>
                    </select>
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">客户级别</label>
                    <select class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                        <option>请选择客户级别</option>
                    </select>
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">地址</label>
                    <select class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                        <option>请选择城市</option>
                    </select>
                </div>
                <div class="col-span-1">
                    <label class="block text-gray-700">详细地址</label>
                    <input type="text" placeholder="请输入详细地址" class="mt-1 block w-full border border-gray-300 rounded-md p-2">
                </div>
                <div class="col-span-2">
                    <label class="block text-gray-700">下次联系时间</label>
                    <div class="flex items-center mt-1">
                        <i class="far fa-clock text-gray-500 mr-2"></i>
                        <input type="text" placeholder="选择下次联系时间" class="block w-full border border-gray-300 rounded-md p-2">
                    </div>
                </div>
                <div class="col-span-2">
                    <label class="block text-gray-700">备注</label>
                    <textarea placeholder="请输入备注" class="mt-1 block w-full border border-gray-300 rounded-md p-2"></textarea>
                </div>
            </div>
            <div class="flex justify-end mt-6">
                <button type="submit" class="bg-blue-500 text-white px-4 py-2 rounded-md mr-2">确定</button>
                <button type="button" class="bg-gray-300 text-gray-700 px-4 py-2 rounded-md">取消</button>
            </div>
        </form>
    </div>
    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css">
</body>
</html>
'''
# ret = ragUI.rag(html)
#print(ret)